package simple.utils.tracking.storage;


import com.google.common.collect.ImmutableMap;
import com.google.protobuf.InvalidProtocolBufferException;
import com.netflix.astyanax.AstyanaxContext;
import com.netflix.astyanax.Keyspace;
import com.netflix.astyanax.MutationBatch;
import com.netflix.astyanax.connectionpool.NodeDiscoveryType;
import com.netflix.astyanax.connectionpool.OperationResult;
import com.netflix.astyanax.connectionpool.exceptions.ConnectionException;
import com.netflix.astyanax.connectionpool.impl.ConnectionPoolConfigurationImpl;
import com.netflix.astyanax.connectionpool.impl.ConnectionPoolType;
import com.netflix.astyanax.impl.AstyanaxConfigurationImpl;
import com.netflix.astyanax.model.Column;
import com.netflix.astyanax.model.ColumnFamily;
import com.netflix.astyanax.model.ColumnList;
import com.netflix.astyanax.thrift.ThriftFamilyFactory;
import org.apache.commons.lang.ArrayUtils;
import org.joda.time.DateTime;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import simple.utils.data.Bytes;
import simple.utils.data.IterableKeyValue;
import simple.utils.data.SeekingIterator;
import simple.utils.tracking.TrackingProto;
import simple.utils.tracking.TrackingUtils;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantLock;


/**
 * Wrapper for Cassandra DB
 * Created by martin on 12.5.14.
 */
public class CassandraStorageFactory implements KeyValueStorageFactory, AutoCloseable {

    private static final Logger LOGGER = LoggerFactory.getLogger(CassandraStorageFactory.class);

    private static final String DEFAULT_KEYSPACE = "tracking";
    private static final byte[] META_PREFIX = new byte[1];
    private static final byte[] DATA_PREFIX = new byte[]{1};
    private static final byte[] META_DAY_TS = "day_ts".getBytes(StandardCharsets.UTF_8);
    // CF contains one meta row with object's day timestamps (for every object one column) and this key accesses it
    private static final CassandraByteArrayKey metaDayTSRowKey = new CassandraByteArrayKey(ArrayUtils.addAll(META_PREFIX, META_DAY_TS));

    private final Map<String, ColumnFamily> cfs = new ConcurrentHashMap<>();
    private final ReentrantLock lock = new ReentrantLock();

    private Keyspace keyspace;
    private AstyanaxContext<Keyspace> ctx;


//    /**
//     * Keyspace defined by name must exists in cassandra db already
//     *
//     * @param keySpaceName keyspace to connect
//     * @throws ConnectionException
//     */
//
//    // todo @Property injection // import atlascon fw
//    public CassandraStorageFactory(String keySpaceName) throws ConnectionException {
//        ctx = new AstyanaxContext.Builder()
//                .forKeyspace(keySpaceName)
//                .withConnectionPoolConfiguration(new ConnectionPoolConfigurationImpl("WhateverNameYouLike")
//                        .setMaxConnsPerHost(100)
//                        .setSeeds("192.168.1.199:9160") // todo
//                )
//                .withAstyanaxConfiguration(
//                        new AstyanaxConfigurationImpl()
//                                .setConnectionPoolType(ConnectionPoolType.ROUND_ROBIN)
//                                .setDiscoveryType(NodeDiscoveryType.NONE)
//                )
//                .buildKeyspace(ThriftFamilyFactory.getInstance());
//        keyspace = ctx.getClient();
//        ctx.start();
//
//        keyspace.dropKeyspace(); // drops and recreates keyspace programmatically, for debugging purposes
//        keyspace.createKeyspace(ImmutableMap.<String, Object>builder()
//                        .put("strategy_options", ImmutableMap.<String, Object>builder()
//                                .put("replication_factor", "1")
//                                .build())
//                        .put("strategy_class", "SimpleStrategy")
//                        .build()
//        );
//    }

    private CassandraStorageFactory() {
    }

    @Override
    public IterableKeyValue<byte[], TrackingProto.TrackingRecord> get(String name) {

        ColumnFamily cf = cfs.get(name);
        lock.lock();
        try {
            if (cf == null) {
                cf = new ColumnFamily<>(name, new CassandraByteArraySerializer(), new CassandraByteArraySerializer());
                cfs.put(name, cf);
                keyspace.createColumnFamily(cf, null);
            }
        } catch (ConnectionException e) {
            LOGGER.error(e.getLocalizedMessage()); // usually if column already exists in cassandra db, but is not registered in cf map
        } finally {
            lock.unlock();
        }
        final ColumnFamily finCf = cf;

        return new IterableKeyValue<byte[], TrackingProto.TrackingRecord>() {

            @Override
            public TrackingProto.TrackingRecord get(byte[] bytes) throws IOException {
                throw new UnsupportedOperationException("Use iterator instead");
            }

            @Override
            public void put(byte[] bytes, TrackingProto.TrackingRecord trackingRecord) throws IOException {
                byte[] dataRowKey = TrackingUtils.idPart(bytes);
                byte[] dataColKey = Arrays.copyOfRange(bytes, dataRowKey.length, bytes.length);

                MutationBatch m = keyspace.prepareMutationBatch();
                CassandraByteArrayKey metaColKey = new CassandraByteArrayKey(dataRowKey);

                DateTime ts = new DateTime(Bytes.bytesToLong(Arrays.copyOf(dataColKey, TrackingUtils.LONG_LENGTH))); // using provided time
                long dayTs = ts.getMillis() - ts.millisOfDay().get();

                try { // update user day timestamp (meta info)
                    OperationResult rs = keyspace.prepareQuery(finCf).getKey(metaDayTSRowKey).execute();
                    ColumnList<CassandraByteArrayKey> list = (ColumnList<CassandraByteArrayKey>) rs.getResult();
                    Column<CassandraByteArrayKey> col = list.getColumnByName(metaColKey);
                    TrackingProto.LongList.Builder longList = TrackingProto.LongList.newBuilder();

                    if (col == null) { // no id registered, no records
                        longList.addLong(dayTs);
                    } else {
                        byte[] protoRawList = col.getByteArrayValue();
                        longList.mergeFrom(protoRawList);
                        if (!longList.getLongList().contains(dayTs)) {
                            longList.addLong(dayTs);
                        }
                    }
                    m.withRow(finCf, metaDayTSRowKey).putColumn(metaColKey, longList.build().toByteArray());
                } catch (ConnectionException e) {
                    throw new IOException(e);
                }

                dataRowKey = ArrayUtils.addAll(DATA_PREFIX, dataRowKey);
                dataRowKey = ArrayUtils.addAll(dataRowKey, Bytes.longToBytes(dayTs));
                m.withRow(finCf, new CassandraByteArrayKey(dataRowKey)).putColumn(new CassandraByteArrayKey(dataColKey), trackingRecord.toByteArray());
                try {
                    m.execute();
                } catch (ConnectionException e) {
                    throw new IOException(e);
                }
            }

            @Override
            public void delete(byte[] bytes) throws IOException {
                throw new UnsupportedOperationException();
            }

            @Override
            public long records() throws IOException {
                SeekingIterator<byte[], TrackingProto.TrackingRecord> it = iterator();
                long i = 0;
                for (; it.hasNext(); it.next(), i++) ;
                return i;
            }

            @Override
            public SeekingIterator<byte[], TrackingProto.TrackingRecord> iterator() {
                return new SeekingIterator<byte[], TrackingProto.TrackingRecord>() {
                    Iterator<Column<CassandraByteArrayKey>> it;
                    List<Long> tss = null;
                    Iterator<Long> tssIt;
                    byte[] originalKey;
                    byte[] plainRowKey;

                    @Override
                    public void seek(byte[] key) {
                        plainRowKey = TrackingUtils.idPart(key);
                        CassandraByteArrayKey metaColKey = new CassandraByteArrayKey(plainRowKey);
                        byte[] tsB = Arrays.copyOfRange(key, plainRowKey.length, plainRowKey.length + TrackingUtils.LONG_LENGTH);
                        long ts = Bytes.bytesToLong(tsB);
                        plainRowKey = ArrayUtils.addAll(DATA_PREFIX, plainRowKey);
//                        System.out.println("Querying: " + Arrays.toString(plainRowKey));

                        try {
                            // reading all object day timestamps
                            OperationResult rs = keyspace.prepareQuery(finCf).getKey(metaDayTSRowKey).execute();
                            ColumnList<CassandraByteArrayKey> list = (ColumnList<CassandraByteArrayKey>) rs.getResult();
                            Column<CassandraByteArrayKey> col = list.getColumnByName(metaColKey);

                            if (col != null) {
                                byte[] protoRawList = col.getByteArrayValue();
                                tss = TrackingProto.LongList.parseFrom(protoRawList).getLongList();
                                tssIt = tss.iterator();
                            }

                            // set seeking iterator to relevant first position
                            if (ts != 0 && tssIt.hasNext()) {
                                long dayTs = tssIt.next();
                                while (ts > dayTs && tssIt.hasNext()) // seek to proper day
                                    dayTs = tssIt.next();

                                byte[] relKey = ArrayUtils.addAll(plainRowKey, Bytes.longToBytes(dayTs)); // one user day
                                rs = keyspace.prepareQuery(finCf).getKey(new CassandraByteArrayKey(relKey)).execute();
                                it = ((ColumnList<CassandraByteArrayKey>) rs.getResult()).iterator();

                                while (it.hasNext()) {
                                    byte[] colB = it.next().getName().getKey();
                                    long colTs = Bytes.bytesToLong(Arrays.copyOf(colB, TrackingUtils.LONG_LENGTH));
                                    if (ts < colTs) // seek to proper record
                                        break;
                                }
                            }

                        } catch (ConnectionException | InvalidProtocolBufferException e) {
                            LOGGER.error(e.getLocalizedMessage());
                        }
                    }

                    @Override
                    public void close() throws Exception {

                    }

                    @Override
                    public boolean hasNext() {
                        if (it != null && it.hasNext()) {
                            return true;
                        } else {
                            if (tssIt != null && tssIt.hasNext()) {
                                Long ts = tssIt.next();
                                byte[] key = ArrayUtils.addAll(plainRowKey, Bytes.longToBytes(ts));
                                OperationResult rs = null;
                                try {
                                    rs = keyspace.prepareQuery(finCf).getKey(new CassandraByteArrayKey(key)).execute();
                                } catch (ConnectionException e) {
                                    LOGGER.error(e.getLocalizedMessage());
                                }
                                Object result = rs.getResult();
                                if (result == null)
                                    return false;
                                it = ((ColumnList<CassandraByteArrayKey>) result).iterator();
                                return it.hasNext();
                            } else {
                                return false;
                            }
                        }
                    }

                    @Override
                    public Map.Entry<byte[], TrackingProto.TrackingRecord> next() {
                        return new Map.Entry<byte[], TrackingProto.TrackingRecord>() {
                            final Column<CassandraByteArrayKey> next = it.next();
                            public TrackingProto.TrackingRecord tr;

                            @Override
                            public byte[] getKey() { // [1,0,0,0,5,103,101,110,73,100,0,0,1,70,55,-120,42,-85,-17,-43]
                                originalKey = ArrayUtils.addAll(plainRowKey, next.getName().getKey());
                                originalKey = Arrays.copyOfRange(originalKey, 1, originalKey.length); // removing first byte (Data | Meta tag)
                                return originalKey;
                            }

                            @Override
                            public TrackingProto.TrackingRecord getValue() {
                                try {
                                    tr = TrackingProto.TrackingRecord.parseFrom(next.getByteArrayValue());
                                    return TrackingProto.TrackingRecord.parseFrom(next.getByteArrayValue());
                                } catch (InvalidProtocolBufferException e) {
                                    LOGGER.error(e.getLocalizedMessage());
                                }

                                return TrackingProto.TrackingRecord.getDefaultInstance();
                            }

                            @Override
                            public TrackingProto.TrackingRecord setValue(TrackingProto.TrackingRecord value) {
                                return value;
                            }

                            @Override
                            public boolean equals(Object o) {
                                if (o == null)
                                    return false;
                                return tr.equals(o);
                            }

                            @Override
                            public int hashCode() {
                                if (originalKey == null)
                                    return 0;
                                return Arrays.hashCode(originalKey) + tr.hashCode() * 37;
                            }
                        };
                    }

                    @Override
                    public void remove() {
                        it.remove();
                    }
                };
            }

            @Override
            public void close() throws Exception {
                // nothing
            }
        };
    }

    @Override
    public void close() throws Exception {
        ctx.shutdown();
        ctx = null;
        keyspace = null;
    }


    public static class CassandraStorageFactoryBuilder {
        private Keyspace keyspace;
        private AstyanaxContext<Keyspace> ctx;
        private CassandraStorageFactory csf;
        private String keySpaceName = DEFAULT_KEYSPACE;

        private boolean withDefaultConfig;
        private String seeds = "192.168.1.199:9160"; // default

        public static CassandraStorageFactoryBuilder newBuilder() {
            return new CassandraStorageFactoryBuilder();
        }

        private CassandraStorageFactoryBuilder() {
            csf = new CassandraStorageFactory();
        }

        public CassandraStorageFactoryBuilder withDefaultConfig() {
            this.withDefaultConfig = true;
            return this;
        }

        public CassandraStorageFactoryBuilder withAstyanaxContext(AstyanaxContext<Keyspace> ctx) {
            this.ctx = ctx;
            return this;
        }

        public CassandraStorageFactoryBuilder withKeySpaceName(String keySpaceName) {
            this.keySpaceName = keySpaceName;
            return this;
        }

        public CassandraStorageFactoryBuilder withSeeds(String seeds) {
            this.seeds = seeds;
            return this;
        }

        public CassandraStorageFactory build() {
            if (ctx == null || withDefaultConfig) {
                csf.ctx = new AstyanaxContext.Builder()
                        .forKeyspace(keySpaceName)
                        .withConnectionPoolConfiguration(new ConnectionPoolConfigurationImpl("WhateverNameYouLike")
                                        .setMaxConnsPerHost(100)
                                        .setSeeds(seeds)
                        )
                        .withAstyanaxConfiguration(
                                new AstyanaxConfigurationImpl()
                                        .setConnectionPoolType(ConnectionPoolType.ROUND_ROBIN)
                                        .setDiscoveryType(NodeDiscoveryType.NONE)
                        )
                        .buildKeyspace(ThriftFamilyFactory.getInstance());
            }

            csf.keyspace = csf.ctx.getClient();
            csf.ctx.start();

            try {
                csf.keyspace.createKeyspace(ImmutableMap.<String, Object>builder()
                                .put("strategy_options", ImmutableMap.<String, Object>builder()
                                        .put("replication_factor", "1")
                                        .build())
                                .put("strategy_class", "SimpleStrategy")
                                .build()
                );
            } catch (ConnectionException e) {
                LOGGER.error(e.getLocalizedMessage());
                e.printStackTrace();
            }
            return csf;
        }


    }


}
